from django.core.management.base import BaseCommand, CommandError

from host.models import Host, City, Room


class Command(BaseCommand):
    help = '统计主机数量'

    def handle(self, *args, **options):
        self.stdout.write('开始统计主机数量')
        for city in City.objects.all():
            all_count = 0
            for room in city.room_set.all():
                host_count = room.host_set.count()
                room.host_count = host_count
                room.save()
                self.stdout.write(f'{room}主机数量为{host_count}')
                all_count += host_count
            city.host_count = all_count
            city.save()
            self.stdout.write(f'{city}主机数量为{host_count}')
        self.stdout.write('统计完成')
